import numpy as np
from matplotlib import pyplot as plt

from src.straight_line import image_projection, rotate_image, hough_transform, detect_scale_bar, image_preprocessing



def auto_select_best_roi(gray, win_shape=(128, 128), step=32):
    """
    自动选出周期性最显著的小窗口
    :param gray: 原图灰度
    :param win_shape: (h, w) 窗口大小
    :param step: 滑动步长
    :return: (x, y, w, h) ROI 矩形
    """
    H, W = gray.shape
    win_h, win_w = win_shape
    best_score, best_xy = -1, (0, 0)

    # 1. 构造滑动窗口
    ys = range(0, H - win_h + 1, step)
    xs = range(0, W - win_w + 1, step)

    for y in ys:
        for x in xs:
            patch = gray[y:y + win_h, x:x + win_w]

            # 2. FFT + 中心化
            f = np.fft.fft2(patch)
            fshift = np.fft.fftshift(f)
            mag = np.abs(fshift)

            # 3. 环形掩码（去掉直流 & 高频噪声）
            cy, cx = win_h // 2, win_w // 2
            Y, X = np.ogrid[:win_h, :win_w]
            dist = np.sqrt((X - cx) ** 2 + (Y - cy) ** 2)
            mask = (dist > 10) & (dist < min(cx, cy) * 0.4)
            energy = np.sum(mag[mask])

            # 4. 记录最大能量窗口
            if energy > best_score:
                best_score = energy
                best_xy = (x, y)

    x, y = best_xy
    return x, y, win_w, win_h

def main(image_path):
    img, enhanced, edges = image_preprocessing(image_path=image_path)

    # ---- 自动比例尺 ----
    scale_bar_actual, scale_bar_pixels = detect_scale_bar(img, fallback_nm=10.0)

    # ---- 自动 ROI ----
    roi = auto_select_best_roi(img, win_shape=(256, 256), step=64)
    x, y, w, h = roi
    roi_img = img[y:y + h, x:x + w]
    roi_edges = edges[y:y + h, x:x + w]
    print(f"🎯 自动选定 ROI: x={x}, y={y}, w={w}, h={h}")

    # ---- 后续保持不变 ----
    rotation_angle, roi_edges_uint8 = hough_transform(roi_edges=roi_edges, roi_img=roi_img)
    rotated_roi, rotated_edges = rotate_image(roi_img=roi_img,
                                              roi_edges_uint8=roi_edges_uint8,
                                              rotation_angle=rotation_angle)
    image_projection(rotated_roi=rotated_roi,
                     scale_bar_actual=scale_bar_actual,
                     scale_bar_pixels=scale_bar_pixels)

if __name__ == "__main__":
    image_path = "../img/new/0006.tif"  # 图像路径
    plt.rc('font', family='Microsoft YaHei')
    main(image_path)